from pylab import *  # The plotting module
import random
import copy

gen1000 = list()

def genertions1000():
    pop_list = ["A"] * 50
    pop_list.extend(["B"] * 50)

    gen1000.append(copy.deepcopy(pop_list)) # init generation

    tmpGen = list()
    for i in range(1000):
        for j in range(100):
            tmpGen.append(gen1000[i][random.randint(0, 99)])

        gen1000.append(copy.deepcopy(tmpGen))

        # check if 'A' disappears
        if not "A" in tmpGen:
            print("A disappears at generation of :", i)
            return

        # check if 'B' disappears
        if not "B" in tmpGen:
            print("B disappears at generation of :", i)
            return

        tmpGen.clear()

# ================= #

def countGenAB(gen):
    cnt_A = 0
    cnt_B = 0
    for i in gen:
        if i == "A":
            cnt_A += 1
        elif i == "B":
            cnt_B += 1

    return (cnt_A, cnt_B)

# ======== main ========= #
genAcntLst = list()
genBcntLst = list()

genertions1000();

for i in gen1000:
    cntA, cntB = countGenAB(i)
    genAcntLst.append(cntA)
    genBcntLst.append(cntB)

x_data = list()
for i in range(len(gen1000)):
    x_data.append(i)

plot(x_data, genAcntLst, 'r', label="gen A count")   # Plot the data. The 'r' sets the colour to red
plot(x_data, genBcntLst, 'b', label="gen B count")   # Plot the data. The 'b' sets the colour to blue
xlabel("genertion #")
ylabel("gen A / B count")
title("alleles AB population")
legend(loc='best')  # lower right

show()  # Display the plot
